#include "graph/operator.h"
#include "graph/operator_reg.h"


namespace ge {
#define _THIS_TYPE GroupedMatmul

#define DYNAMIC_INPUT_FUNCS(x) \
TensorDesc get_dynamic_input_desc_##x(uint32_t index) const {                                                        \
    return Operator::GetDynamicInputDesc(#x, index);                                                                   \
  }                                                                                                                    \
  graphStatus update_dynamic_input_desc_##x(uint32_t index, const TensorDesc &tensorDesc) {                            \
    return Operator::UpdateDynamicInputDesc(#x, index, tensorDesc);                                                    \
  }                                                                                                                    \
  _THIS_TYPE &set_dynamic_input_##x(uint32_t dstIndex, Operator &v) {                                                  \
    Operator::SetInput(#x, dstIndex, v);                                                                               \
    return *this;                                                                                                      \
  }                                                                                                                    \
  ATTRIBUTED_DEPRECATED(_THIS_TYPE &set_dynamic_input_##x(uint32_t, Operator &, const char *))                         \
  _THIS_TYPE &set_dynamic_input_##x(uint32_t dstIndex, Operator &v, const std::string &srcName) {                      \
    Operator::SetInput(#x, dstIndex, v, srcName.c_str());                                                              \
    return *this;                                                                                                      \
  }                                                                                                                    \
  _THIS_TYPE &set_dynamic_input_##x(uint32_t dstIndex, Operator &v, const char *srcName) {                             \
    Operator::SetInput(#x, dstIndex, v, srcName);                                                                      \
    return *this;                                                                                                      \
  }

#define DYNAMIC_OUTPUT_FUNCS(x) \
TensorDesc get_dynamic_output_desc_##x(uint32_t index) const {                                                       \
    return Operator::GetDynamicOutputDesc(#x, index);                                                                  \
  }                                                                                                                    \
  graphStatus update_dynamic_output_desc_##x(uint32_t index, const TensorDesc &tensorDesc) {                           \
    return Operator::UpdateDynamicOutputDesc(#x, index, tensorDesc);                                                   \
  }

#define OPTIONAL_INPUT_FUNCS(x) \
_THIS_TYPE &set_input_##x(Operator &v) {                                                                             \
    Operator::SetInput(#x, v);                                                                                         \
    return *this;                                                                                                      \
  }                                                                                                                    \
  ATTRIBUTED_DEPRECATED(_THIS_TYPE &set_input_##x##_by_name(Operator &, const char *))                                 \
  _THIS_TYPE &set_input_##x(Operator &v, const std::string &srcName) {                                                 \
    Operator::SetInput(#x, v, srcName.c_str());                                                                        \
    return *this;                                                                                                      \
  }                                                                                                                    \
  _THIS_TYPE &set_input_##x##_by_name(Operator &v, const char *srcName) {                                              \
    Operator::SetInput(#x, v, srcName);                                                                                \
    return *this;                                                                                                      \
  }                                                                                                                    \
  _THIS_TYPE &set_input_##x(Operator &v, uint32_t index) {                                                             \
    Operator::SetInput(#x, v, index);                                                                                  \
    return *this;                                                                                                      \
  }                                                                                                                    \
  TensorDesc get_input_desc_##x() const {                                                                              \
    return Operator::GetInputDescByName(#x);                                                                           \
  }                                                                                                                    \
  graphStatus update_input_desc_##x(const TensorDesc &tensorDesc) {                                                    \
    return Operator::UpdateInputDesc(#x, tensorDesc);                                                                  \
  }

#define ATTR_FUNCS(x, Type) \
_THIS_TYPE &set_attr_##x(const Op##Type &v) {                                                                        \
  Operator::SetAttr(#x, v);                                                                                          \
  return *this;                                                                                                      \
}                                                                                                                    \
_THIS_TYPE &set_attr_##x(const function<Op##Type()> &v) {                                                            \
  (void) v;                                                                                                          \
  return *this;                                                                                                      \
}

class GroupedMatmul: public Operator {
 public:
  GroupedMatmul(const char *name, 
    int64_t x_num=0, 
    int64_t weight_num=0, 
    int64_t bias_num=0, 
    int64_t scale_num=0,
    int64_t offset_num=0,
    int64_t antiquant_scale_num=0,
    int64_t antiquant_offset_num=0,
    int64_t y_num=0) : Operator(name, "GroupedMatmul") {
    DynamicInputRegister("x", x_num);
    DynamicInputRegister("weight", weight_num);
    DynamicInputRegister("bias", bias_num);
    DynamicInputRegister("scale", scale_num);
    DynamicInputRegister("offset", offset_num);
    DynamicInputRegister("antiquant_scale", antiquant_scale_num);
    DynamicInputRegister("antiquant_offset", antiquant_offset_num);
    OptionalInputRegister("group_list");
    OptionalInputRegister("per_token_scale");
    DynamicOutputRegister("y", y_num);

    AttrRegister("split_item", static_cast<int64_t>(0));
    AttrRegister("dtype", static_cast<int64_t>(0));
    AttrRegister("transpose_weight", false);
    AttrRegister("transpose_x", false);
    AttrRegister("group_type", static_cast<int64_t>(-1));
    AttrRegister("group_list_type", static_cast<int64_t>(0));
    AttrRegister("act_type", static_cast<int64_t>(0));
    std::vector<int64_t> act_typtuning_confige = {0};
    AttrRegister("act_typtuning_confige", act_typtuning_confige);
  }

  DYNAMIC_INPUT_FUNCS(x)
  DYNAMIC_INPUT_FUNCS(weight)
  DYNAMIC_INPUT_FUNCS(bias)
  DYNAMIC_INPUT_FUNCS(scale)
  DYNAMIC_INPUT_FUNCS(offset)
  DYNAMIC_INPUT_FUNCS(antiquant_scale)
  DYNAMIC_INPUT_FUNCS(antiquant_offset)

  OPTIONAL_INPUT_FUNCS(group_list)
  OPTIONAL_INPUT_FUNCS(per_token_scale)

  DYNAMIC_OUTPUT_FUNCS(y)

  ATTR_FUNCS(split_item, Int)
  ATTR_FUNCS(dtype, Int)
  ATTR_FUNCS(transpose_weight, Bool)
  ATTR_FUNCS(transpose_x, Bool)
  ATTR_FUNCS(group_type, Int)
  ATTR_FUNCS(group_list_type, Int)
  ATTR_FUNCS(act_type, Int)
  ATTR_FUNCS(tuning_config, ListInt)
};
#undef _THIS_TYPE
#undef DYNAMIC_INPUT_FUNCS
#undef DYNAMIC_OUTPUT_FUNCS
#undef OPTIONAL_INPUT_FUNCS
#undef ATTR_FUNCS
}